#!/usr/bin/env python
#-*- coding: utf-8 -*-

import numpy as np
import math
import scipy
import matplotlib.pyplot as plt

import my_utils as mu

import vehicles.fixed_wing.dynamic_model_3d as dm
import vehicles.fixed_wing.control_3d       as ctl
import vehicles.fixed_wing.plot_3d          as mp

np.set_printoptions(precision=2, suppress=False, linewidth=200)
# loads and displays the dynamic model's parameters
#dyn_name = "../config/Rcam_single_engine.xml"
dyn_name = "../config/epicea_ref.xml"
#dyn_name = "../config/epicea_wing_moved.xml"
p_dm = dm.Param(dyn_name)
print p_dm

def run_sim(X0, time):
    X = np.zeros((time.size,dm.s_size))
    U = np.zeros((len(time), p_dm.input_nb))
    X[0] = X0
    for i in range(1,time.size):
        U[i-1] = ap.run(time[i-1], X[i-1,:])
        foo,X[i,:] = scipy.integrate.odeint(dm.dyn, X[i-1,:], [time[i-1], time[i]], args=(U[i-1], p_dm, ))
    U[-1] = U[-2]
    return X, U

# run simulations
time = np.arange(0., 10.0, 0.01)

ap = ctl.Ap(p_dm)
ap.set_mode(2)
debug = False


def run_pid_sim():
    X0 = np.array(ap.Xe); X0[dm.s_phi] += mu.rad_of_deg(-10.)

    K = np.zeros((p_dm.input_nb, dm.s_size))
    ap.set_gain(K)
    ap._print()
    Xol, Uol = run_sim(X0, time)

    K = np.zeros((p_dm.input_nb, dm.s_size))
    K[dm.i_da1+p_dm.eng_nb, dm.s_phi] = 4.
    K[dm.i_da1+p_dm.eng_nb, dm.s_p] = 0.5
    ap.set_gain(K)
    ap._print()

    Xcl, Ucl = run_sim(X0, time)
    fig1 = mp.plot_dynamics(time, Xol, window_title="{:s} ({:s})".format("State",p_dm.name))
    fig1 = mp.plot_dynamics(time, Xcl, figure=fig1, legend=["open loop", "corrected"])
    fig2 = mp.plot_command(time, Uol, figure=None, window_title="Command", legend=None, filename=None)
    fig2 = mp.plot_command(time, Ucl, figure=fig2)


def run_pid_sim2():
    X0 = np.array(ap.Xe);
    X0[dm.s_phi]   += mu.rad_of_deg(-10.)
    X0[dm.s_theta] += mu.rad_of_deg(-10.)
    X0[dm.s_beta]  += mu.rad_of_deg(-10.)

    K = np.zeros((p_dm.input_nb, dm.s_size))
    K[dm.i_da1+p_dm.eng_nb, dm.s_phi] = 4.
    K[dm.i_da1+p_dm.eng_nb, dm.s_p] = 0.5
    K[dm.i_de1+p_dm.eng_nb, dm.s_theta] = 2.
    K[dm.i_de1+p_dm.eng_nb, dm.s_q] = 2.5
    K[dm.i_dr1+p_dm.eng_nb, dm.s_beta] = -2.5
    K[dm.i_dr1+p_dm.eng_nb, dm.s_r] = 1.5
    
    ap.set_gain(K)
    ap._print()

    Xcl, Ucl = run_sim(X0, time)
    fig1 = mp.plot_dynamics(time, Xcl, window_title="{:s} ({:s})".format("State",p_dm.name))
    fig2 = mp.plot_command(time, Ucl, figure=None, window_title="Command", legend=None, filename=None)



def run_longi_sim_placement():
    """Longitudinal"""
    X0 = np.array(ap.Xe); X0[dm.s_theta] += mu.rad_of_deg(-10.)
    ap.set_spec(omega_pitch=None, omega_roll=None, xi=0.8, debug=debug)
    Xol, Uol = run_sim(X0, time)
    omegas = [1., 1.5, 3., 4.]; Trajs = []
    legend = ["Natural"]+["$\omega = {:.1f}$".format(o) for o in omegas]
    for om in omegas:
        ap.set_spec(omega_pitch=om, omega_roll=None, xi=0.8, debug=debug)
        Trajs.append(run_sim(X0, time))
    fig1 = mp.plot_longitudinal_dynamics(time, Xol, Uol, window_title="{:s} ({:s})".format("State",p_dm.name))
    for Xcl, Ucl in Trajs:
        fig1 = mp.plot_longitudinal_dynamics(time, Xcl, Ucl, figure=fig1, legend=legend)


def run_lat_sim_placement():
    """Lateral"""
    X0 = np.array(ap.Xe); X0[dm.s_phi] += mu.rad_of_deg(10.)
    ap.set_spec(omega_pitch=None, omega_roll=None, debug=debug)
    ap.get_mode()._print(ap)
    Xol, Uol = run_sim(X0, time)
    omegas = [1., 2., 4., 8.]; Trajs = []
    for om in omegas:
        ap.set_spec(omega_pitch=None, omega_roll=om, xi=2.5, debug=False)
        ap._print()
        Trajs.append(run_sim(X0, time))

#    print legend
    fig1 = mp.plot_dynamics(time, Xol, window_title="{:s} ({:s})".format("State",p_dm.name))
    fig2 = mp.plot_command(time, Uol, figure=None, window_title="Command", legend=None, filename=None)
    for Xcl, Ucl in Trajs:
        fig1 = mp.plot_dynamics(time, Xcl, figure=fig1)
        fig2 = mp.plot_command(time, Ucl, figure=fig2)


import control.matlab
def run_sim_lqr():
    time = np.arange(0., 30.0, 0.01)
    X0 = np.array(ap.Xe);
    X0[dm.s_phi]   += mu.rad_of_deg(-10.)
    X0[dm.s_theta] += mu.rad_of_deg(-10.)
    X0[dm.s_psi] += mu.rad_of_deg(-10.)
    X0[dm.s_beta]  += mu.rad_of_deg(10.)

    A, B = ap.get_jacobian()
    A1=A[dm.s_alpha:, dm.s_alpha:]
    B1=B[dm.s_alpha:,1:]
    #            alpha beta   phi  theta  psi     p       q       r
    Q = np.diag([0.5,   0.5,  0.5, 2.5,   2.,  0.0005, 0.0005, 0.0005])
    #            ail  ele   rud
    R = np.diag([0.1, 0.5, 0.1])
    (K, X1, E) = control.matlab.lqr(A1, B1, Q, R)
    K1 = np.zeros((p_dm.input_nb, dm.s_size))
    K1[ap.P_dm.eng_nb:,dm.s_alpha:] = -K
    ap.set_gain(K1)
    ap._print()
    Xcl, Ucl = run_sim(X0, time)
    fig1 = mp.plot_dynamics(time, Xcl, window_title="{:s} ({:s})".format("State",p_dm.name))
    fig2 = mp.plot_command(time, Ucl, figure=None, window_title="Command", legend=None, filename=None)


def run_sim_lqr2():
    time = np.arange(0., 30.0, 0.01)
    X0 = np.array(ap.Xe);
    X0[dm.s_phi]   += mu.rad_of_deg(-10.)
    X0[dm.s_theta] += mu.rad_of_deg(-10.)
    X0[dm.s_psi] += mu.rad_of_deg(-10.)
    X0[dm.s_beta]  += mu.rad_of_deg(10.)

    A, B = ap.get_jacobian()
    A1=A[dm.s_v:, dm.s_v:]
    B1=B[dm.s_v:,:]
    #             v   alpha beta   phi  theta  psi     p       q       r
    Q = np.diag([0.5, 0.5,   0.5,  0.5, 2.5,   2.,  0.0005, 0.0005, 0.0005])
    #            eng  ail  ele   rud
    R = np.diag([0.2, 0.1, 0.5, 0.1])
    (K, X1, E) = control.matlab.lqr(A1, B1, Q, R)
    K1 = np.zeros((p_dm.input_nb, dm.s_size))
    K1[:,dm.s_v:] = -K
    ap.set_gain(K1)
    ap._print()
    Xcl, Ucl = run_sim(X0, time)
    fig1 = mp.plot_dynamics(time, Xcl, window_title="{:s} ({:s})".format("State",p_dm.name))
    fig2 = mp.plot_command(time, Ucl, figure=None, window_title="Command", legend=None, filename=None)

def run_sim_lqr3():
    time = np.arange(0., 30.05, 0.01)
    X0 = np.array(ap.Xe);
    X0[dm.s_phi]   += mu.rad_of_deg(-10.)
    X0[dm.s_theta] += mu.rad_of_deg(-10.)
    X0[dm.s_psi] += mu.rad_of_deg(-10.)
    X0[dm.s_beta]  += mu.rad_of_deg(10.)

    A, B = ap.get_jacobian()
    A1=A[dm.s_z:, dm.s_z:]
    B1=B[dm.s_z:,:]
    #             z        v     alpha beta   phi  theta  psi     p       q       r
    Q = np.diag([0.01,  0.01,  0.5,   0.5,  0.5, 2.5,   2.,  0.0005, 0.0005, 0.0005])
    #            eng  ail  ele   rud
    R = np.diag([0.3, 0.1, 0.7, 0.1])
    (K, X1, E) = control.matlab.lqr(A1, B1, Q, R)
    K1 = np.zeros((p_dm.input_nb, dm.s_size))
    K1[:,dm.s_z:] = -K
    ap.set_gain(K1)
    ap._print()
    Xcl, Ucl = run_sim(X0, time)
    fig1 = mp.plot_dynamics(time, Xcl, window_title="{:s} ({:s})".format("State",p_dm.name))
    fig2 = mp.plot_command(time, Ucl, figure=None, window_title="Command", legend=None, filename=None)
        
    
#run_pid_sim()    
#run_pid_sim2()
#run_longi_sim_placement()
#run_lat_sim_placement()
#run_sim_lqr()
#run_sim_lqr2()
run_sim_lqr3()

plt.show()

#if __name__ == "__main__":


    
